{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Using Decorators to Create Context Managers using Generator Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the last video we saw how we could create a generic class that could take a generator function that had a specific structure and turn it into a context manager.\n",
    "\n",
    "Let's see if we can do one step better, using a decorator instead."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall the basic structure our generator function needs to have:\n",
    "\n",
    "```\n",
    "def gen(args):\n",
    "    # set up happens here, or inside try\n",
    "    try:\n",
    "        yield obj # whatever normally gets returned by __enter__\n",
    "    finally:\n",
    "        # perform clean up code here\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First let's define a generator function to open a file, yield it, and then close it - same as the example we saw in the previous video:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def open_file(fname, mode='r'):\n",
    "    print('opening file...')\n",
    "    f = open(fname, mode)\n",
    "    try:\n",
    "        yield f\n",
    "    finally:\n",
    "        print('closing file...')\n",
    "        f.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let's re-create the context manager wrapper we did in the last video, but this time, I'm going to pass it a generator object, instead of a generator function. But basically the same idea:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class GenContextManager:\n",
    "    def __init__(self, gen):\n",
    "        self.gen = gen\n",
    "        \n",
    "    def __enter__(self):\n",
    "        return next(self.gen)\n",
    "        \n",
    "    def __exit__(self, exc_type, exc_value, exc_tb):\n",
    "        print('calling next to perform cleanup in generator')\n",
    "        try:\n",
    "            next(self.gen)\n",
    "        except StopIteration:\n",
    "            pass\n",
    "        return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At this point we can use this with our generator function as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file...\n",
      "calling next to perform cleanup in generator\n",
      "closing file...\n"
     ]
    }
   ],
   "source": [
    "file_gen = open_file('test.txt', 'w')\n",
    "\n",
    "with GenContextManager(file_gen) as f:\n",
    "    f.writelines('Sir Spamalot')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can read back from the file too:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file...\n",
      "['Sir Spamalot']\n",
      "calling next to perform cleanup in generator\n",
      "closing file...\n"
     ]
    }
   ],
   "source": [
    "file_gen = open_file('test.txt')\n",
    "with GenContextManager(file_gen) as f:\n",
    "    print(f.readlines())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Of course, our context manager object is not very robust - there is no exception handling for example. But let's leave that \"minor detail\" aside for now :-)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We still have to create the generator object from the generator function before we can use the context manager class.\n",
    "\n",
    "We can simplify things even more by using a decorator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def context_manager_dec(gen_fn):\n",
    "    def helper(*args, **kwargs):\n",
    "        gen = gen_fn(*args, **kwargs)\n",
    "        ctx = GenContextManager(gen)\n",
    "        return ctx\n",
    "    return helper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice what this decorator does.\n",
    "\n",
    "It decorates a generator function and returns `helper`. When we invoke `helper` it will create an instance of the generator, and create and return an instance of the context manager.\n",
    "\n",
    "Let's try it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@context_manager_dec\n",
    "def open_file(fname, mode='r'):\n",
    "    print('opening file...')\n",
    "    f = open(fname, mode)\n",
    "    try:\n",
    "        yield f\n",
    "    finally:\n",
    "        print('closing file...')\n",
    "        f.close()    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file...\n",
      "['Sir Spamalot']\n",
      "calling next to perform cleanup in generator\n",
      "closing file...\n"
     ]
    }
   ],
   "source": [
    "with open_file('test.txt') as f:\n",
    "    print(f.readlines())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So now we have an approach to using a decorator to turn any generator function (that has the structure we mentioned earlier) into a context manager!\n",
    "\n",
    "Our code was not very robust, either in the context manager class or in the decorator - and it would take quite a bit more work to make it so.\n",
    "\n",
    "Fortunately the standard library already has this implemented for us - in fact that was one of the critical goals of Python's context managers - the ability to create context managers using generator functions (see PEP 343)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from contextlib import contextmanager"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def open_file(fname, mode='r'):\n",
    "    print('opening file...')\n",
    "    f = open(fname, mode)\n",
    "    try:\n",
    "        yield f\n",
    "    finally:\n",
    "        print('closing file...')\n",
    "        f.close() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "opening file...\n",
      "['Sir Spamalot']\n",
      "closing file...\n"
     ]
    }
   ],
   "source": [
    "with open_file('test.txt') as f:\n",
    "    print(f.readlines())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And of course, this works for more than just opening and closing files. \n",
    "\n",
    "Here are some more examples:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Example 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's implement a timer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from time import perf_counter, sleep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def timer():\n",
    "    stats = dict()\n",
    "    start = perf_counter()\n",
    "    stats['start'] = start\n",
    "    yield stats\n",
    "    end = perf_counter()\n",
    "    stats['end'] = end\n",
    "    stats['elapsed'] = end - start"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with timer() as stats:\n",
    "    sleep(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'start': 2.708953062782967e-07, 'end': 1.0097690265350079, 'elapsed': 1.0097687556397017}\n"
     ]
    }
   ],
   "source": [
    "print(stats)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Example 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, let's redirect `stdout`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "@contextmanager\n",
    "def out_to_file(fname):\n",
    "    current_stdout = sys.stdout\n",
    "    file = open(fname, 'w')\n",
    "    sys.stdout = file\n",
    "    try:\n",
    "        yield None\n",
    "    finally:\n",
    "        file.close()\n",
    "        sys.stdout = current_stdout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with out_to_file('test.txt'):\n",
    "    print('line 1')\n",
    "    print('line 2')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['line 1\\n', 'line 2\\n']\n"
     ]
    }
   ],
   "source": [
    "with open('test.txt') as f:\n",
    "    print(f.readlines())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And of course, `stdout` is back to \"normal\":"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "line 1\n"
     ]
    }
   ],
   "source": [
    "print('line 1')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `contextlib` module actually implements a `stdout` redirect context manager, so we technically don't have to write one ourselves.\n",
    "\n",
    "The difference from the one we wrote is that it needs an open file object, not just a file name. So we would have to open the file, then redirect stdout. We can do this easily by nesting two context managers as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from contextlib import redirect_stdout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "with open('test.txt', 'w') as f:\n",
    "    with redirect_stdout(f):\n",
    "        print('Look on the bright side of life')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And we can check that this worked:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['Look on the bright side of life\\n']\n"
     ]
    }
   ],
   "source": [
    "with open('test.txt') as f:\n",
    "    print(f.readlines())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
